package com.lysky.selfllm.controller;

import cn.hutool.extra.pinyin.PinyinUtil;
import cn.hutool.http.HttpUtil;
import com.lysky.selfllm.service.AIWriter;
import dev.langchain4j.agent.tool.P;
import dev.langchain4j.agent.tool.Tool;
import dev.langchain4j.agent.tool.ToolSpecification;
import dev.langchain4j.agent.tool.ToolSpecifications;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentByLineSplitter;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ToolExecutionResultMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.request.ChatRequestParameters;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.content.injector.ContentInjector;
import dev.langchain4j.rag.content.injector.DefaultContentInjector;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.tool.DefaultToolExecutor;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.net.URISyntaxException;
import java.nio.charset.StandardCharsets;
import java.time.LocalDateTime;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicReference;

@Slf4j
@RestController
public class ChatController {

    @Resource
    private ChatLanguageModel chatLanguageModel;

    @Resource
    private ChatMemory chatMemory;

    @Resource
    private EmbeddingModel embeddingModel;

    @Resource
    private EmbeddingStore<TextSegment> embeddingStore;

    @GetMapping("/chat")
    public String chat(@RequestParam("message") String message) throws Exception {
        chatMemory.add(UserMessage.from(message));
        ToolSpecification timeTool = ToolSpecifications.toolSpecificationFrom(ChatController.class.getMethod("nowTime", null));
        ToolSpecification weatherTool = ToolSpecifications.toolSpecificationFrom(ChatController.class.getMethod("weather", String.class));
        ChatRequest chatRequest = ChatRequest.builder()
                .messages(chatMemory.messages())
                .parameters(ChatRequestParameters.builder().toolSpecifications(timeTool, weatherTool).build())
                .build();
        AtomicReference<ChatResponse> chatResponse = new AtomicReference<>(chatLanguageModel.chat(chatRequest));
        AiMessage aiMessage = chatResponse.get().aiMessage();
        if (aiMessage.hasToolExecutionRequests()){
            aiMessage.toolExecutionRequests().forEach(request -> {
                DefaultToolExecutor defaultToolExecutor = new DefaultToolExecutor(this, request);
                String result = defaultToolExecutor.execute(request, UUID.randomUUID().toString());
                log.warn("执行结果：{}", result);
                ToolExecutionResultMessage toolExecutionResultMessage = ToolExecutionResultMessage.from(request, result);
                chatMemory.add(toolExecutionResultMessage);
            });
            chatResponse.set(chatLanguageModel.chat(chatMemory.messages()));
        } else {
            chatMemory.add(chatResponse.get().aiMessage());
        }
        return chatResponse.get().aiMessage().text();
    }

    @GetMapping("/write")
    public String write(@RequestParam("content") String content) {
        AIWriter writer = AiServices.create(AIWriter.class, chatLanguageModel);
        return writer.write(content);
    }

    @GetMapping("/design")
    public String design(@RequestParam("name") String name,@RequestParam("attack") String attack,@RequestParam("defense") String defense, @RequestParam("speed") String speed) {
        AIWriter writer = AiServices.builder(AIWriter.class)
                .chatLanguageModel(chatLanguageModel)
                .chatMemory(chatMemory)
//                .chatMemoryProvider(memoryId -> chatMemory)
                .build();
        return writer.design(name, attack, defense, speed);
    }

    @GetMapping("/embed")
    public String embed(@RequestParam("content") String content) {
        TextSegment textSegment = TextSegment.textSegment(content);
        Response<Embedding> embed = embeddingModel.embed(textSegment);
        embeddingStore.add(embed.content(), textSegment);
        return embed.content().toString();
    }

    @GetMapping("/compare")
    public String compare(@RequestParam("content") String content) {
        Response<Embedding> embed = embeddingModel.embed(content);// 找出最相似的embedding
        EmbeddingSearchRequest request = EmbeddingSearchRequest.builder()
                .queryEmbedding(embed.content())
                .maxResults(3)
                .build();
        EmbeddingSearchResult<TextSegment> search = embeddingStore.search(request);
        List<EmbeddingMatch<TextSegment>> matches = search.matches();
        for (EmbeddingMatch<TextSegment> match : matches) {
            log.warn("{}相似度：{}",match.embedded().text() , match.score());
        }
        return matches.toString();
    }

    @GetMapping("/read")
    public String read() throws URISyntaxException {
        // 读取本地知识文件
        Document document = FileSystemDocumentLoader.loadDocument("/Users/feathermoon/Desktop/lysky/JavaProject/self-llm/src/main/resources/knowledge.txt", new TextDocumentParser(StandardCharsets.UTF_8));
        // 把知识文件分解成一个一个的知识条目
        DocumentSplitter documentSplitter = new DocumentByLineSplitter(50, 10);
        List<TextSegment> textSegments = documentSplitter.split(document);
        // 对每个条目进行文本向量化，并存储到embedding store(opensearch)中
        Response<List<Embedding>> response = embeddingModel.embedAll(textSegments);
        embeddingStore.addAll(response.content(), textSegments);
        return "喂饭成功!";
    }

    @GetMapping("/ask")
    public String ask(@RequestParam("question") String question) {
        // 将问题进行向量化，并与embedding store中的文本进行匹配，找出最相似的文本条目
        ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .maxResults(5)
                .minScore(0.4)
                .build();
        Query query = Query.from(question);
        List<Content> contents = contentRetriever.retrieve(query);
        // 将问题和相关知识条目整合，发送给大模型
        ContentInjector contentInjector = new DefaultContentInjector();
        ChatMessage injected = contentInjector.inject(contents, UserMessage.from(question));
        ChatResponse chatResponse = chatLanguageModel.chat(injected);
        return chatResponse.aiMessage().text();
    }










    @Tool("获取当前日期")
    public String nowTime(){
        return LocalDateTime.now().toString();
    }

    @Tool("获取某一个城市今天的天气情况")
    public String weather(@P("指定的城市") String city){
        // 发送天气查询请求
        String url = "https://api.seniverse.com/v3/weather/now.json?key=Sjtio-OGcHsoKmtOb&location=" + transferToPinYin(city) +"&language=zh-Hans&unit=c";
        String s = HttpUtil.get(url);
        log.warn("天气查询URL：{}", url);
        log.warn("天气查询结果：{}", s);
        return s;
    }

    protected String transferToPinYin(String chinese){
        // todo: 实现将中文汉字转为拼音形式，首字母大写
        String pinyin = PinyinUtil.getPinyin(chinese).replace(" ", "");
        return pinyin.substring(0, 1).toUpperCase() + pinyin.substring(1);
    }
}
